You write custom CUDA kernels to replace the PyTorch operators in the given EvoNorm architecture to get speedups.
You have complete freedom to choose the set of operators you want to replace. You may make the decision to replace some operators with custom CUDA kernels and leave others unchanged. You may replace multiple operators with custom implementations, consider operator fusion opportunities (combining multiple operators into a single kernel, for example, combining normalization+affine_transform+nonlinear_gating), or algorithmic changes (such as optimized memory access patterns). You are only limited by your imagination.

Poisson Negative Log-Likelihood Loss CUDA Optimization with Fused Kernel
CUDA Optimization Techniques
1. Parallel Reduction Architecture
Grid-Stride Loop Pattern: Each thread processes multiple elements with stride gridDim.x * blockDim.x

Block-Level Reduction: Partial sums computed in shared memory

Dynamic Grid Sizing: grid_size = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE

2. Fused Kernel Design
Single Kernel Execution: Combines all loss computations in one kernel launch

Branch Handling: Efficiently handles log_input and full mode conditions

Mathematical Fusion: Integrates exponential, logarithmic, and conditional operations

3. Memory Access Optimization
Coalesced Memory Access: Sequential reading of input and target tensors

Shared Memory Utilization: s_data[BLOCK_SIZE] for block-level sum reduction

Contiguous Tensors: Ensures input and target are contiguous in memory

4. Numerical Stability Features
Epsilon Protection: eps_val prevents log(0) in non-log mode

Stirling's Approximation: Conditional Stirling term for full mode when y > 1.0f

Float Safety: Proper handling of edge cases and special values

5. Mathematical Operations
7. Performance Optimizations
Minimal Global Memory Writes: Only thread 0 writes block sum to global memory

Efficient Thread Utilization: All threads participate in computation and reduction

Load Balancing: Grid-stride loops handle arbitrary tensor sizes

Constant Propagation: PI and configuration parameters as compile-time constants

8. Implementation Features
Configuration Flexibility: Supports log_input, full, and eps parameters

Comprehensive Validation: Tensor device, contiguity, and size checking

Edge Case Handling: Empty tensor detection and proper zero handling

PyTorch Integration: Seamless tensor passing and automatic differentiation support

Key CUDA Concepts Used
Grid-Stride Loops for workload distribution across all elements

Shared Memory Reduction for parallel sum computation

Conditional Execution for handling different mathematical modes

Memory Coalescing for efficient global memory access

Kernel Fusion combining multiple mathematical operations

Workflow Summary
Configuration Setup: Parse log_input, full, and eps parameters

Memory Preparation: Ensure contiguous tensor layouts

Kernel Launch: Execute fused Poisson NLL computation with parallel reduction

Final Reduction: Sum block partial sums and compute mean loss

Mathematical Components
Exponential Computation: expf() for log-input mode

Logarithmic Computation: logf() for non-log mode and Stirling term

Stirling's Approximation: Complete term for Poisson distribution normalization

Element-wise Operations: Parallel computation across all tensor elements

Expected Performance Benefits
2-4x speedup over PyTorch implementation for large tensors

Reduced kernel launches through operation fusion

Better memory efficiency through coalesced access patterns

Scalable performance with increasing tensor sizes

This implementation provides a production-ready Poisson NLL Loss with significant performance improvements through careful CUDA optimization, parallel reduction patterns, and numerical stability considerations.


Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:
import torch
import torch.nn as nn
import torch.nn.functional as F

BATCH_SIZE = 4096
FEATURE_DIM = 512

# --- 损失函数的参数 ---
# log_input=True: loss = exp(input) - target * input
# log_input=False: loss = input - target * log(input + eps)
LOG_INPUT = True
# full=True: 添加 Stirling's approximation
FULL = False
EPS = 1e-8


class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.criterion = nn.PoissonNLLLoss(
            log_input=LOG_INPUT,
            full=FULL,
            eps=EPS,
            reduction='mean'
        )

    def forward(self, input_tensor: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        # 在 PyTorch 中，input_tensor 是文档中的 'input'
        return self.criterion(input_tensor, target)


def get_inputs():
    # Input (log_input=True 时) 可以是任意实数
    input_tensor = torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32)

    # Target 在 Poisson 分布中代表计数，且在 'full' 模式下会计算 log(target)
    # 因此 target 必须是 >= 0 的。我们使用 rand 来确保
    target = torch.rand(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32) * 10  # 乘以 10 以便有一些 > 1

    return [input_tensor, target]


def get_init_inputs():
    return []